import json
import sys
import argparse
import subprocess
import os
import yaml
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


def check_ssh_access(ip):
    try:
        if sys.version_info >= (3, 5):
            subprocess.run(['ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=5', ip, 'echo', 'SSH_OK'], check=True,
                           stdout=subprocess.DEVNULL)
        else:
            subprocess.check_call(['ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=5', ip, 'echo', 'SSH_OK'],
                                  stdout=subprocess.DEVNULL)
            # process = subprocess.Popen(['ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=5', ip, 'echo', 'SSH_OK'], stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
            # process.communicate()
            # if process.returncode != 0:
            #     raise subprocess.CalledProcessError(process.returncode, process.args)
        return True
    except subprocess.CalledProcessError as e:
        logging.error("SSH connection to %s failed: %s", ip, e)
        return False


def get_device_ips(ip, master_ip):
    device_ips = []
    command_prefix = ['hccn_tool', '-i', '-ip', '-g']
    try:
        for i in range(8):
            command_prefix.insert(2, str(i))
            if ip != master_ip:
                command_prefix.insert(0, 'ssh')
                command_prefix.insert(1, ip)
            if sys.version_info >= (3, 5):
                result = subprocess.run(command_prefix, stdout=subprocess.PIPE, check=True)
                stdout = result.stdout
            else:
                stdout = subprocess.check_output(command_prefix, stderr=subprocess.STDOUT)
                # process = subprocess.Popen(command_prefix, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
                # stdout, stderr = process.communicate()
                # if process.returncode != 0:
                #     raise subprocess.CalledProcessError(process.returncode, command_prefix, stderr)
            output = stdout.decode('utf-8').strip()
            ipaddr = output.split('\n')[0].split(':')[1].strip()
            device_ips.append(ipaddr)
            if ip != master_ip:
                command_prefix.pop(0)
                command_prefix.pop(0)
            command_prefix.pop(2)
        return device_ips
    except subprocess.CalledProcessError as e:
        logging.error("Failed to get device IPs from %s: %s", ip, e.output.decode('utf-8'))
        return []


def setup_ssh_keys(server_ips, master_ip, local_public_key_path):
    for ip in server_ips:
        if ip != master_ip and not check_ssh_access(ip):
            try:
                command = [
                    "ssh-copy-id",
                    "-o", "StrictHostKeyChecking=no",
                    "-o", "UserKnownHostsFile=/dev/null",
                    "-i", local_public_key_path,
                    "root@%s" % ip
                ]
                if sys.version_info >= (3, 5):
                    subprocess.run(command, check=True)
                else:
                    subprocess.check_output(command, stderr=subprocess.STDOUT)
                if not check_ssh_access(ip):
                    logging.error("%s 配置免密失败", ip)
                    return False
            except subprocess.CalledProcessError as e:
                logging.error("Failed to copy SSH key to %s: %s", ip, e.output.decode('utf-8'))
                return False
    return True


def update_json(json_file, master_ip, server_id, container_ip, server_count, server_index, all_servers):
    if os.path.exists(json_file):
        with open(json_file, 'r') as f:
            data = json.load(f)
    else:
        data = {
            "version": "1.0",
            "server_count": str(server_count),
            "server_list": [],
            "status": "completed"
        }

    if data.get("server_count") == str(server_count):
        pass
    else:
        data["server_count"] = str(server_count)

    rank_id_start = (server_index - 1) * 8
    server_found = False
    devices = [{"device_id": str(i), "device_ip": ip, "rank_id": str(rank_id_start + i)} for i, ip in
               enumerate(get_device_ips(server_id, master_ip))]

    for server in data["server_list"]:
        if server["server_id"] == server_id:
            server_found = True
            if container_ip:
                server["container_ip"] = container_ip
            server["device"] = devices
            break

    if not server_found:
        new_server = {
            "server_id": server_id,
            "container_ip": container_ip if container_ip else server_id,
            "device": devices,
            "host_nic_ip": "reserve"
        }
        data["server_list"].append(new_server)

    data["server_list"] = sorted(data["server_list"], key=lambda x: int(x["device"][0]["rank_id"]))

    with open(json_file, 'w') as f:
        json.dump(data, f, indent=4)
    if server_index == server_count:
        os.chmod(json_file, 0o640)
        for ip, _ in all_servers:
            if ip != master_ip:
                try:
                    if sys.version_info >= (3, 5):
                        subprocess.run(["scp", json_file, f"root@{ip}:{json_file}"])
                    else:
                        subprocess.check_output(["scp", json_file, "root@%s:%s" % (ip, json_file)],
                                                stderr=subprocess.STDOUT)
                except subprocess.CalledProcessError as e:
                    logging.error("Failed to distribute file to %s: %s", ip, e.output.decode('utf-8'))


def main():
    parser = argparse.ArgumentParser(description="Update JSON file with server information.")
    parser.add_argument('--json_file', type=str, default='/root/rank_table.json')
    parser.add_argument('--all_config_file', type=str, default='./all_config.yaml')

    args = parser.parse_args()

    with open(args.all_config_file, 'r') as f:
        config = yaml.safe_load(f)
        all_servers = config['all_ip']

    local_private_key_path = '/root/.ssh/id_rsa'
    local_public_key_path = '/root/.ssh/id_rsa.pub'

    if not (os.path.exists(local_private_key_path) and os.path.exists(local_public_key_path)):
        try:
            subprocess.run(["ssh-keygen", "-t", "rsa"])
        except subprocess.CalledProcessError as e:
            logging.error("Failed to generate SSH keys: %s", e.output.decode('utf-8'))

    master_ip = all_servers[0][0]
    beta_pass = setup_ssh_keys([ip for ip, _ in all_servers], master_ip, local_public_key_path)
    if not beta_pass:
        return logging.info("程序终止,请手动配置免密")

    for index, (server_ip, container_ip) in enumerate(all_servers, start=1):
        try:
            update_json(args.json_file, master_ip, server_ip, container_ip, len(all_servers), index, all_servers)
        except Exception as e:
            logging.error("Failed to update JSON for server %s: %s", server_ip, e)


if __name__ == "__main__":
    main()
